import numpy as np
import pygad
import multiprocessing

# 向量化适应度函数
def fitness_func_vectorized(population):
    # 使用NumPy的向量化操作计算适应度值
    fitness = np.sum(population, axis=1)
    return fitness

# 并行适应度函数
def parallel_fitness(population):
    # 将种群分成多个子集进行并行计算
    num_processes = multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=num_processes)

    # 将种群分成多个子集，每个子集由一个进程处理
    sub_populations = np.array_split(population, num_processes)
    fitness_values = pool.map(fitness_func_vectorized, sub_populations)

    # 将各子集的结果合并
    fitness_values = np.concatenate(fitness_values)
    return fitness_values

# 自定义适应度函数包装器以适应 pygad 的需求
def fitness_func_wrapper(ga_instance, solution, solution_idx):
    global global_population, global_population_fitness
    if solution_idx == 0:  # 仅在第一次调用时计算整个种群的适应度
        global_population_fitness = parallel_fitness(global_population)
    return global_population_fitness[solution_idx]

# 基因数量
num_genes = 4

# 遗传算法参数
sol_per_pop = 20  # 种群大小
num_parents_mating = 10
num_generations = 50
crossover_probability = 0.8
mutation_probability = 0.1
gene_space = [0, 1]  # 基因的取值范围

# 初始化全局变量用于存储种群和适应度值
global_population = np.random.randint(0, 2, (sol_per_pop, num_genes))
global_population_fitness = np.zeros(sol_per_pop)

# 初始化遗传算法实例
ga_instance = pygad.GA(
    num_generations=num_generations,
    num_parents_mating=num_parents_mating,
    fitness_func=fitness_func_wrapper,
    sol_per_pop=sol_per_pop,
    num_genes=num_genes,
    gene_space=gene_space,
    crossover_probability=crossover_probability,
    mutation_probability=mutation_probability
)

# 运行遗传算法
ga_instance.run()

# 获取最佳解决方案
solution, solution_fitness, solution_idx = ga_instance.best_solution()
print(f"最佳解决方案: {solution}")
print(f"最佳解决方案的适应度值: {solution_fitness}")